#!/usr/bin/python3
'''
usage: funtrace++ </path/to/clang++ or g++> [-funtrace-instr-thresh=N] [-funtrace-no-trace=<mangled-names-list-file>] [-funtrace-do-trace=<names-file>] [-funtrace-ignore-loops] <compiler command>
'''
import os, sys, subprocess

class TraceSupressor:
    def __init__(self, args):
        self.instr_thresh = 0
        self.ignore_loops = False
        self.no_trace = []
        self.do_trace = []
        self.verbose = False
        def val(arg): return arg.split('=')[-1]
        def lines(file): return file, open(file).read().strip().split()
        for arg in args:
            if arg.startswith('-funtrace-instr-thresh='):
                self.instr_thresh = int(val(arg))
            elif arg.startswith('-funtrace-no-trace='):
                self.no_file, self.no_trace = lines(val(arg))
            elif arg.startswith('-funtrace-do-trace='):
                self.do_file, self.do_trace = lines(val(arg))
            elif arg == '-funtrace-ignore-loops':
                self.ignore_loops = True
            elif arg == '-funtrace-verbose':
                self.verbose = True

    def suppress(self, funcname, num_instr, loops):
        if funcname in self.do_trace:
            if self.verbose:
                print(f'{funcname} listed in the do-trace file {self.do_file}')
            return
        reason = None
        if funcname in self.no_trace:
            reason = f'{funcname} listed in the trace suppression file {self.no_file}'
        elif num_instr < self.instr_thresh:
            if not loops:
                reason = f'{funcname} has {num_instr} instructions, less than -funtrace-instr-thresh={self.instr_thresh}'
            elif self.ignore_loops:
                reason = f'{funcname} has {num_instr} instructions, less than -funtrace-instr-thresh={self.instr_thresh}; it has {loops} loops but -funtrace-ignore-loops was passed'
        if self.verbose and reason:
            print(reason)
        return reason

# note that we don't support filtering XRay's output, in part on the theory that it already has
# -fxray-instruction-threshold=N; to support it, we'd need to look for the NOPs it inserts - it doesn't
# put in call instructions, that's done by runtime code patching
hooks = [
    '__cyg_profile_func_enter',
    '__cyg_profile_func_exit',
    '__fentry__',
    '__return__',
]

def filter_asm(asm_file, suppressor):
    with open(asm_file) as f:
        lines = f.read().split('\n')

    funcname = None
    infunc = False
    instrs = 0
    loops = 0
    labels = []
    funcstart = None

    changed = False

    for i,line in enumerate(lines):
        l = line.strip()
        if l.startswith('.type') and l.endswith('@function'):
            funcname = l.split(',')[0].split()[-1]
        elif l == '.cfi_startproc':
            #print('in func',funcname)
            infunc = True
            funcstart = i+1
            instrs = 0
            loops = 0
            labels = []
        elif l == '.cfi_endproc':
            #print('end func', funcname, instrs, loops)
            infunc = False
            suppression_reason = suppressor.suppress(funcname, instrs, loops)
            if not suppression_reason:
                continue
            
            for j in range(funcstart, i):
                l = lines[j].strip()
                for hook in hooks:
                    if 'call' in l or 'jmp' in l:
                        if hook in l:
                            lines[j] = '# ' + lines[j] + ' # ' + suppression_reason
                            if 'jmp' in l: # tail call
                                lines[j] = '    ret ' + lines[j]
                            changed = True
                            break
        elif infunc:
            if not l:
                continue
            t = l.split()[0]
            isinstr = line[0].isspace() and not t.startswith('.') and not t.endswith(':')
            if isinstr:
                instrs += 1
                for label in labels:
                    if label in l:
                        loops += 1
                        break
            elif t.startswith('.') and t.endswith(':'):
                labels.append(t[:-1])

    if changed:
        with open(asm_file, 'w') as f:
            f.write('\n'.join(lines))

def exec_compiler(cmd, execl=True):
    compiler = cmd[0]
    if not os.path.exists(compiler):
        compiler = subprocess.getoutput(f'which {compiler}')
    if execl:
        os.execl(compiler, *cmd)
    else:
        subprocess.run([compiler]+cmd[1:])

def compile_filter_and_assemble(cmd, funtrace_args):
    suppressor = TraceSupressor(funtrace_args)
    compile_to_asm_cmd, assemble_cmd, asm_file = compile_and_assemble_commands(cmd)

    #print(' '.join(compile_to_asm_cmd))
    exec_compiler(compile_to_asm_cmd, execl=False)

    filter_asm(asm_file, suppressor)

    #print(' '.join(assemble_cmd))
    exec_compiler(assemble_cmd, execl=False)

def compile_and_assemble_commands(cmd):
    ofile = None
    cfile = None
    sfile = None

    extensions = 'c cpp cc cxx cp CPP c++ C'.split()
    def is_src_arg(arg):
        if arg.startswith('-'):
           return False
        for ext in extensions:
           if arg.endswith(ext):
               return True

    for i,arg in enumerate(cmd):
        if arg == '-o' and i+1 < len(cmd):
            ofile = cmd[i+1]
            sfile = ofile+'.s'
        elif is_src_arg(arg):
            cfile = arg

    if cfile:
        if not ofile:
            ofile = cfile[:cfile.rfind('.')] + '.o'
            sfile = cfile[:cfile.rfind('.')] + '.s'
            compile_to_asm_cmd = [('-S' if arg == '-c' else arg) for arg in cmd] + ['-o',sfile]
        else:
            compile_to_asm_cmd = [('-S' if arg == '-c' else (sfile if arg == ofile else arg)) for arg in cmd]
    else:
        print(f'funtrace++ - WARNING: -c passed but could not determine the input source file in `{cmd}`')
        exec_compiler(cmd)

    assemble_cmd = [cmd[0], '-c', sfile, '-o', ofile]
    if 'clang' in cmd[0]:
        assemble_cmd += ['-Wa,-W'] # clang produces assembly using MD5 sums for some source files but not others and then the assembler warns of
        # "inconsistent use of md5 sums", not sure how to suppress this better...

    return compile_to_asm_cmd, assemble_cmd, sfile

def main():
    cmd = sys.argv[1:]
    funtrace_args = [arg for arg in cmd if arg.startswith('-funtrace-')]
    cmd = [arg for arg in cmd if not arg.startswith('-funtrace-')]
    if '-c' in cmd and '-E' not in cmd and '-dM' not in cmd and funtrace_args:
        compile_filter_and_assemble(cmd, funtrace_args)
    else:
        exec_compiler(cmd)

if __name__ == '__main__':
    main()
